/**
 * 2025/01/01 - 2025/12/31
 */
package com.snail.starter.db.holder;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.transaction.support.TransactionCallback;

import java.util.function.Consumer;

/**
 * @author zgadmin
 * @date 2025/8/1 16:27
 */
@Slf4j
@Component
public class DynamicDataSourceTemplate {

    /**
     * 在当前数据源下执行, 不抛异常
     *
     * @param dataSourceFunction 业务执行逻辑
     */
    public static void execute(DynamicDataSourceFunction dataSourceFunction) {
        try {
            executeWithThrowable(dataSourceFunction);
        } catch (Throwable throwable) {
            log.error("", throwable);
            log.error("execute Throwable default datasource");
        }
    }

    /**
     * 在指定数据源下执行, 不抛异常
     *
     * @param dataSourceId       指定数据源ID
     * @param dataSourceFunction 业务执行逻辑
     */
    public static void execute(String dataSourceId, DynamicDataSourceFunction dataSourceFunction) {
        try {
            executeWithThrowable(dataSourceId, dataSourceFunction);
        } catch (Throwable throwable) {
            log.error("", throwable);
            log.error("execute Throwable data-source:{}", dataSourceId);
        }
    }

    public static void executeWithThrowable(DynamicDataSourceFunction dataSourceFunction) {
        dataSourceFunction.execute();
    }

    public static void executeWithThrowable(String dataSourceId, DynamicDataSourceFunction dataSourceFunction) {
        try {
            if (null != dataSourceId && !"".equals(dataSourceId)) {
                log.info("execute change data-source,from:{},to:{}", DynamicDataSourceContextHolder.peek(), dataSourceId);
                DynamicDataSourceContextHolder.push(dataSourceId);
            }
            executeWithThrowable(dataSourceFunction);
        } finally {
            if (null != dataSourceId && !"".equals(dataSourceId)) {
                DynamicDataSourceContextHolder.poll();
                log.info("execute change data-source finish,form:{},to:{}", dataSourceId, DynamicDataSourceContextHolder.peek());
            }
        }
    }

    /**
     * 当前数据源下操作,需要返回,没有异常
     *
     * @param dataSourceFunction 事务控制逻辑
     * @param <T>                返回结果的类型
     * @return 返回结果
     */
    public static <T> T callback(DynamicDataSourceCallBackFunction<T> dataSourceFunction) {
        try {
            return callbackWithThrowable(dataSourceFunction);
        } catch (Throwable throwable) {
            log.error("", throwable);
            log.error("callback Throwable default datasource");
            return null;
        }
    }

    /**
     * 指定数据源下操作,需要返回值,没有异常
     *
     * @param dataSourceId       数据源ID
     * @param dataSourceFunction 事务控制逻辑
     * @param <T>                返回结果的类型
     * @return 返回结果
     */
    public static <T> T callback(String dataSourceId, DynamicDataSourceCallBackFunction<T> dataSourceFunction) {
        try {
            return callbackWithThrowable(dataSourceId, dataSourceFunction);
        } catch (Throwable throwable) {
            log.error("", throwable);
            log.error("callback Throwable data-source:{}", dataSourceId);
            return null;
        }
    }

    /**
     * 当前数据源下操作,需要返回,有异常
     *
     * @param dataSourceFunction 事务控制逻辑
     * @param <T>                返回结果的类型
     * @return 返回结果
     */
    public static <T> T callbackWithThrowable(DynamicDataSourceCallBackFunction<T> dataSourceFunction) {
        return dataSourceFunction.callback();
    }

    /**
     * 指定数据源下操作,需要返回值,有异常
     *
     * @param dataSourceId       数据源ID
     * @param dataSourceFunction 事务控制逻辑
     * @param <T>                返回结果的类型
     * @return 返回结果
     */
    public static <T> T callbackWithThrowable(String dataSourceId, DynamicDataSourceCallBackFunction<T> dataSourceFunction) {
        try {
            if (null != dataSourceId && !"".equals(dataSourceId)) {
                log.info("callbackWithThrowable change data-source,from:{},to:{}", DynamicDataSourceContextHolder.peek(), dataSourceId);
                DynamicDataSourceContextHolder.push(dataSourceId);
            }
            return callbackWithThrowable(dataSourceFunction);
        } finally {
            if (null != dataSourceId && !"".equals(dataSourceId)) {
                DynamicDataSourceContextHolder.poll();
                log.info("callbackWithThrowable change data-source finish,form:{},to:{}", dataSourceId, DynamicDataSourceContextHolder.peek());
            }
        }
    }

    public interface DynamicDataSourceCallBackFunction<T> {

        T callback();

    }

    public interface DynamicDataSourceFunction {

        void execute();

    }

    public interface DsTransactionCallback<T> extends TransactionCallback<T> {

    }

    public interface DsTransactionConsumer<T> extends Consumer<T> {

    }

}